#pragma once
#include "opencv2/core/core.hpp"

class RectPnPSolver
{
public:
    RectPnPSolver(const cv::Mat &camera_matrix, const cv::Mat &dist_coeff,
                  double target_width = 0, double target_height = 0)
    {
        camera_matrix.copyTo(cam_matrix);
        dist_coeff.copyTo(distortion_coeff);
        width_target = target_width;
        height_target = target_height;
    }

    void setTargetSize(double width, double height)
    {
        width_target = width;
        height_target = height;
    }

    void setCameraParam(const cv::Mat & camera_matrix, const cv::Mat &dist_coeff)
    {
        camera_matrix.copyTo(cam_matrix);
        dist_coeff.copyTo(distortion_coeff);
    }

    void solvePnP4Points(const std::vector<cv::Point2f> &points2d, cv::Mat &rot, cv::Mat &trans);

public:
    cv::Mat cam_matrix;
    cv::Mat distortion_coeff;
    double width_target;
    double height_target;
};

class AngleSolver:public RectPnPSolver
{
public:
    AngleSolver(const cv::Mat & camera_matrix, const cv::Mat & dist_coeff,
                double target_width = 0, double target_height = 0, double z_scale = 1.0,
                double min_dist = 50.0, double max_dist = 600.0)
        :RectPnPSolver(camera_matrix, dist_coeff, target_width, target_height)
    {
        min_distance = min_dist;
        max_distance = max_dist;

        rot_camera2ptz = cv::Mat::eye(3,3,CV_64FC1);
        trans_camera2ptz = cv::Mat::zeros(3,1,CV_64FC1);
        offset_y_barrel_ptz = 0;
        scale_z = z_scale;
    }

    void setScaleZ(double scale){scale_z = scale;}
    void setRelationPoseCameraPTZ(const cv::Mat & rot_camera_ptz, const cv::Mat & trans_camera_ptz, double y_offset_barrel_ptz);
    void getTarget2dPointstion(const cv::RotatedRect &rect, std::vector<cv::Point2f> &target2d, const cv::Point2f & offset);

    bool getAngle(const cv::RotatedRect &rect, double &angle_x, double& angle_y, double bullet_speed = 0,
                  double current_ptz_angle = 0.0, const cv::Point2f &offset = cv::Point2f());

    void tranformationCamera2PTZ(const cv::Mat &pos, cv::Mat &transed_pos);

    void adjustPTZ2Barrel(const cv::Mat & pos_in_ptz, double &angle_x, double &angle_y, double bullet_speed = 0.0, double current_ptz_angle = 0.0);

public:
    cv::Mat position_in_camera;
    cv::Mat position_in_ptz;
private:
    cv::Mat trans_camera2ptz;
    cv::Mat rot_camera2ptz;

    // offset between barrel and ptz on y axis (cm)
    double offset_y_barrel_ptz;

    // scope of detection distance (cm)
    double min_distance;
    double max_distance;
    double scale_z;

};

class AngleSolverFactory
{
public:
    AngleSolverFactory(AngleSolver * angle_solver = nullptr):solver(angle_solver){}

    typedef enum {TARGET_BUFF, TARGET_ARMOR, TARGET_SMALL_ARMOR} TargetType;

    void setSolver(AngleSolver *angle_solver){solver = angle_solver;}

    AngleSolver & getSolver(){return *solver;}

    void setTargetSize(double width, double height, TargetType type);

    void adjustRect2FixedRatio(cv::RotatedRect &rect, double wh_ratio)
    {
        rect.size.height = rect.size.width / wh_ratio;
    }
    bool getAngle(const cv::RotatedRect &rect, TargetType type, double &angle_x, double &angle_y,
                  double bullet_speed, double current_ptz_angle, const cv::Point2f &offset = cv::Point2f());

private:
    double armor_width;
    double armor_height;
    double small_armor_width;
    double small_armor_height;
    double buff_width;
    double buff_height;
    AngleSolver * solver;
};
